package com.yungu.swift.admin.handler;

import com.google.common.base.Joiner;
import com.google.common.util.concurrent.RateLimiter;
import com.yungu.swift.base.ResponseData;
import com.yungu.swift.utils.JsonUtils;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.CollectionUtils;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.config.BeanPostProcessor;
import org.springframework.stereotype.Component;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.handler.HandlerInterceptorAdapter;
import org.springframework.web.servlet.mvc.condition.PatternsRequestCondition;
import org.springframework.web.servlet.mvc.method.RequestMappingInfo;
import org.springframework.web.servlet.mvc.method.annotation.RequestMappingHandlerMapping;
import org.springframework.web.util.UrlPathHelper;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;

import static com.yungu.swift.admin.consts.AdminErrorEnum.SERVER_BUSY;

/**
 * 限流处理器 —— 基于令牌桶方案
 *
 * @author : cuixiuyin
 * @date : 2019/8/1
 */
@Slf4j
@Component
public class RateLimitHandler extends HandlerInterceptorAdapter implements BeanPostProcessor {

    /**
     * 默认全局每秒并发 10 个请求
     */
    private static final Integer GLOBAL_RATE_LIMITER = 10;
    /**
     * 每个 url 的具体限流方案
     */
    public static Map<PatternsRequestCondition, RateLimiter> URL_RATE_MAP;


    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        response.setContentType("application/json;charset=UTF-8");
        if (URL_RATE_MAP != null) {
            String lookupPath = new UrlPathHelper().getLookupPathForRequest(request);
            for (PatternsRequestCondition patternsRequestCondition : URL_RATE_MAP.keySet()) {
                //使用 Spring DispatcherServlet 的匹配器 PatternsRequestCondition 进行匹配
                List<String> matches = patternsRequestCondition.getMatchingPatterns(lookupPath);
                if (CollectionUtils.isEmpty(matches)) {
                    continue;
                }
                // 允许2秒内的突发等待
                if (!URL_RATE_MAP.get(patternsRequestCondition).tryAcquire(2000, TimeUnit.MILLISECONDS)) {
                    log.warn(" 请求 '{}' 匹配到 {}，超过限流速率，获取令牌失败。", lookupPath, Joiner.on(",").join(patternsRequestCondition.getPatterns()));
                    // 服务器繁忙，降级处理
                    response.getWriter().print(JsonUtils.writeJsonStr(ResponseData.buildErrorResponse(SERVER_BUSY.getCode(), SERVER_BUSY.getMsg())));
                    return false;
                }
                log.info(" 请求 '{}' 匹配到 {} ,成功获取令牌，进入请求。", lookupPath, Joiner.on(",").join(patternsRequestCondition.getPatterns()));
                break;
            }
        }
        return super.preHandle(request, response, handler);
    }

    @Override
    public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException {
        return bean;
    }

    @Override
    public Object postProcessAfterInitialization(Object bean, String beanName) throws BeansException {
        if (RequestMappingHandlerMapping.class.isAssignableFrom(bean.getClass())) {
            if (URL_RATE_MAP == null) {
                URL_RATE_MAP = new ConcurrentHashMap<>(16);
            }
            log.info("we get all the controllers's methods and assign it to urlRateMap");
            RequestMappingHandlerMapping requestMappingHandlerMapping = (RequestMappingHandlerMapping) bean;
            Map<RequestMappingInfo, HandlerMethod> handlerMethods = requestMappingHandlerMapping.getHandlerMethods();
            for (RequestMappingInfo mappingInfo : handlerMethods.keySet()) {
                PatternsRequestCondition requestCondition = mappingInfo.getPatternsCondition();
                URL_RATE_MAP.put(requestCondition, RateLimiter.create(GLOBAL_RATE_LIMITER));
            }
        }
        return bean;
    }
}